"""
BCI Competition IV 2a数据集GDF文件解析脚本
使用mne-python库读取和解析.gdf文件
"""

import os
import numpy as np
import mne
from pathlib import Path


def parse_gdf_file(file_path):
    """
    解析单个GDF文件
    
    参数:
        file_path: GDF文件路径
    
    返回:
        raw: mne Raw对象，包含原始数据
        info: 文件信息字典
    """
    print(f"\n{'='*60}")
    print(f"正在解析文件: {os.path.basename(file_path)}")
    print(f"{'='*60}")
    
    # 读取GDF文件
    try:
        raw = mne.io.read_raw_gdf(file_path, preload=True, verbose=False)
    except Exception as e:
        print(f"读取文件时出错: {e}")
        return None, None
    
    # 获取基本信息
    info = {
        'file_path': file_path,
        'n_channels': len(raw.ch_names),
        'n_samples': raw.n_times,
        'sampling_freq': raw.info['sfreq'],
        'duration': raw.n_times / raw.info['sfreq'],
        'channel_names': raw.ch_names,
        'data_shape': raw.get_data().shape
    }
    
    # 打印基本信息
    print(f"\n基本信息:")
    print(f"  通道数: {info['n_channels']}")
    print(f"  采样点数: {info['n_samples']}")
    print(f"  采样频率: {info['sampling_freq']} Hz")
    print(f"  持续时间: {info['duration']:.2f} 秒")
    print(f"  数据形状: {info['data_shape']}")
    
    # 打印通道名称
    print(f"\n通道名称 ({len(raw.ch_names)}个):")
    for i, ch_name in enumerate(raw.ch_names):
        print(f"  {i+1:2d}. {ch_name}")
    
    # 获取事件标记（如果存在）
    try:
        events, event_id = mne.events_from_annotations(raw, verbose=False)
        if len(events) > 0:
            print(f"\n事件标记:")
            print(f"  事件数量: {len(events)}")
            print(f"  事件类型: {event_id}")
            
            # 显示前10个事件
            print(f"\n前10个事件:")
            print(f"  时间点(s) | 事件ID | 事件类型")
            print(f"  {'-'*40}")
            for i, event in enumerate(events[:10]):
                time_sec = event[0] / raw.info['sfreq']
                event_type = [k for k, v in event_id.items() if v == event[2]]
                event_type_str = event_type[0] if event_type else "未知"
                print(f"  {time_sec:8.2f} | {event[2]:6d} | {event_type_str}")
            if len(events) > 10:
                print(f"  ... (还有 {len(events) - 10} 个事件)")
        else:
            print(f"\n事件标记: 未找到事件标记")
    except Exception as e:
        print(f"\n事件标记: 无法读取事件信息 ({e})")
    
    # 获取数据统计信息
    data = raw.get_data()
    print(f"\n数据统计信息:")
    print(f"  数据范围: [{data.min():.4f}, {data.max():.4f}]")
    print(f"  均值: {data.mean():.4f}")
    print(f"  标准差: {data.std():.4f}")
    
    return raw, info


def parse_all_gdf_files(folder_path):
    """
    解析文件夹中的所有GDF文件
    
    参数:
        folder_path: 包含GDF文件的文件夹路径
    """
    folder = Path(folder_path)
    gdf_files = list(folder.glob("*.gdf"))
    
    if len(gdf_files) == 0:
        print(f"在 {folder_path} 中未找到GDF文件")
        return
    
    print(f"\n找到 {len(gdf_files)} 个GDF文件")
    print(f"文件列表:")
    for f in sorted(gdf_files):
        print(f"  - {f.name}")
    
    # 解析所有文件
    all_info = []
    for gdf_file in sorted(gdf_files):
        raw, info = parse_gdf_file(gdf_file)
        if info is not None:
            all_info.append(info)
    
    # 打印汇总信息
    if all_info:
        print(f"\n\n{'='*60}")
        print(f"汇总信息")
        print(f"{'='*60}")
        print(f"\n总共解析了 {len(all_info)} 个文件")
        print(f"\n文件统计:")
        print(f"  文件名 | 通道数 | 采样频率 | 持续时间(秒) | 数据形状")
        print(f"  {'-'*70}")
        for info in all_info:
            filename = os.path.basename(info['file_path'])
            print(f"  {filename:15s} | {info['n_channels']:6d} | {info['sampling_freq']:8.1f} | "
                  f"{info['duration']:12.2f} | {info['data_shape']}")


def extract_data_from_gdf(file_path, start_time=None, stop_time=None):
    """
    从GDF文件中提取数据
    
    参数:
        file_path: GDF文件路径
        start_time: 开始时间（秒），默认为None（从开头）
        stop_time: 结束时间（秒），默认为None（到结尾）
    
    返回:
        data: numpy数组，形状为 (n_channels, n_samples)
        channel_names: 通道名称列表
        sampling_freq: 采样频率
        times: 时间点数组
    """
    raw = mne.io.read_raw_gdf(file_path, preload=True, verbose=False)
    
    # 提取指定时间段的数据
    if start_time is not None or stop_time is not None:
        raw = raw.crop(tmin=start_time, tmax=stop_time)
    
    data = raw.get_data()
    channel_names = raw.ch_names
    sampling_freq = raw.info['sfreq']
    times = raw.times
    
    return data, channel_names, sampling_freq, times


if __name__ == "__main__":
    # 获取当前脚本所在目录
    current_dir = Path(__file__).parent
    
    print("BCI Competition IV 2a数据集GDF文件解析工具")
    print("="*60)
    
    # 解析所有文件
    parse_all_gdf_files(current_dir)
    
    # 示例：解析单个文件
    print(f"\n\n{'='*60}")
    print("示例：提取单个文件的数据")
    print(f"{'='*60}")
    
    # 找到第一个文件作为示例
    gdf_files = list(current_dir.glob("*.gdf"))
    if gdf_files:
        example_file = sorted(gdf_files)[0]
        print(f"\n提取文件: {example_file.name}")
        
        # 提取前5秒的数据作为示例
        data, channel_names, sfreq, times = extract_data_from_gdf(
            str(example_file), 
            start_time=0, 
            stop_time=5
        )
        
        print(f"\n提取的数据:")
        print(f"  数据形状: {data.shape}")
        print(f"  时间点数量: {len(times)}")
        print(f"  时间范围: [{times[0]:.2f}, {times[-1]:.2f}] 秒")
        print(f"\n前3个通道的前10个采样点:")
        for i in range(min(3, len(channel_names))):
            print(f"  {channel_names[i]}: {data[i, :10]}")

